import os
import django
import openpyxl
from django.conf import settings
from django.db import transaction
from openpyxl.styles import Alignment
from rest_framework.views import APIView
from rest_framework.response import Response
from rest_framework import status

from aitest.models import AITestCase, Users, AITestDetail
from aitest.resources import AITestCaseResource, AITestDetailResource, BuildInfoResource

class ExportExcelView(APIView):
    def get(self, request, *args, **kwargs):
        # 获取查询参数
        test_summary = request.query_params.get('test_summary', None)
        user_email = request.query_params.get('user', None)

        # 根据查询参数过滤数据
        aitest_case_resource = AITestCaseResource()
        aitest_detail_resource = AITestDetailResource()
        build_info_resource = BuildInfoResource()

        params = {}
        if test_summary:
            params['test_summary'] = test_summary
        if user_email:
            params['user'] = user_email

        aitest_case_dataset = aitest_case_resource.export(aitest_case_resource.get_queryset(**params))
        heads1 = aitest_case_dataset.headers
        values1 = aitest_case_dataset.dict

        aitest_detail_dataset = aitest_detail_resource.export()
        heads2 = aitest_detail_dataset.headers
        values2 = aitest_detail_dataset.dict

        total_heads = heads1 + heads2[1:]  # 排除重复的'id'
        for index, test in enumerate(values1):
            new_detail_list = []
            for detail in values2:
                if test['id'] == detail['ai_test_case']:
                    new_detail_list.append(detail)
            new_detail_list.sort(key=lambda x: x['index'])
            test['ai_test_detail'] = new_detail_list
            values1[index] = test

        # 创建Excel文件
        wb = openpyxl.Workbook()
        ws = wb.active
        ws.title = "AITest Data"

        # 写入表头
        for col_num, header in enumerate(total_heads, 1):
            cell = ws.cell(row=1, column=col_num, value=header)
            cell.border = openpyxl.styles.Border(
                left=openpyxl.styles.Side(style='thin'),
                right=openpyxl.styles.Side(style='thin'),
                top=openpyxl.styles.Side(style='thin'),
                bottom=openpyxl.styles.Side(style='thin')
            )

        # 写入数据
        row_num = 2
        for test in values1:
            if test['ai_test_detail']:
                start_row = row_num
                for detail in test['ai_test_detail']:
                    for col_num, key in enumerate(heads1, 1):
                        cell = ws.cell(row=row_num, column=col_num, value=test[key])
                        cell.border = openpyxl.styles.Border(
                            left=openpyxl.styles.Side(style='thin'),
                            right=openpyxl.styles.Side(style='thin'),
                            top=openpyxl.styles.Side(style='thin'),
                            bottom=openpyxl.styles.Side(style='thin')
                        )
                    for col_num, key in enumerate(heads2[1:], len(heads1) + 1):  # 跳过重复的'id'
                        cell = ws.cell(row=row_num, column=col_num, value=detail[key])
                        cell.border = openpyxl.styles.Border(
                            left=openpyxl.styles.Side(style='thin'),
                            right=openpyxl.styles.Side(style='thin'),
                            top=openpyxl.styles.Side(style='thin'),
                            bottom=openpyxl.styles.Side(style='thin')
                        )
                    row_num += 1
                end_row = row_num - 1
                for col_num in range(1, len(heads1) + 1):
                    ws.merge_cells(start_row=start_row, start_column=col_num, end_row=end_row, end_column=col_num)
                    cell = ws.cell(row=start_row, column=col_num)
                    cell.alignment = Alignment(vertical='top')
                    cell.border = openpyxl.styles.Border(
                        left=openpyxl.styles.Side(style='thin'),
                        right=openpyxl.styles.Side(style='thin'),
                        top=openpyxl.styles.Side(style='thin'),
                        bottom=openpyxl.styles.Side(style='thin')
                    )
            else:
                for col_num, key in enumerate(heads1, 1):
                    cell = ws.cell(row=row_num, column=col_num, value=test[key])
                    cell.border = openpyxl.styles.Border(
                        left=openpyxl.styles.Side(style='thin'),
                        right=openpyxl.styles.Side(style='thin'),
                        top=openpyxl.styles.Side(style='thin'),
                        bottom=openpyxl.styles.Side(style='thin')
                    )
                row_num += 1

        # 保存Excel文件到指定路径
        file_path = os.path.join(settings.MEDIA_ROOT, 'aitest_data.xlsx')
        wb.save(file_path)

        # 返回文件路径
        return Response({'file_path': file_path}, status=status.HTTP_200_OK)

class ImportExcelView(APIView):
    def post(self, request, *args, **kwargs):
        file = request.FILES.get('file')
        if not file:
            return Response({'error': 'No file provided'}, status=status.HTTP_400_BAD_REQUEST)

        try:
            wb = openpyxl.load_workbook(file)
            ws = wb['AITest Data']
        except Exception as e:
            return Response({'error': str(e)}, status=status.HTTP_400_BAD_REQUEST)

        headers = [cell.value for cell in ws[1]]
        data = []
        for row in ws.iter_rows(min_row=2, values_only=True):
            data.append(dict(zip(headers, row)))

        try:
            with transaction.atomic():
                for row in data:
                    if row.get('id'):
                        try:
                            test_case = AITestCase.objects.get(id=row['id'])
                        except AITestCase.DoesNotExist:
                            return Response({'error': f"Test case with id {row['id']} does not exist"}, status=status.HTTP_400_BAD_REQUEST)
                    else:
                        test_case = AITestCase()

                    test_case.test_summary = row.get('test_summary')
                    test_case.test_requirements = row.get('test_requirements')
                    test_case.detail = row.get('detail')
                    test_case.remark = row.get('remark')
                    test_case.reserved_field0 = row.get('reserved_field0')
                    test_case.reserved_field1 = row.get('reserved_field1')
                    test_case.reserved_field2 = row.get('reserved_field2')

                    if row.get('user'):
                        try:
                            user = Users.objects.get(email=row['user'])
                            test_case.user = user
                        except Users.DoesNotExist:
                            return Response({'error': f"User with email {row['user']} does not exist"}, status=status.HTTP_400_BAD_REQUEST)

                    test_case.save()

                    # Handle AITestDetail if present
                    if row.get('ai_test_detail'):
                        AITestDetail.objects.filter(ai_test_case=test_case).delete()
                        for detail in row['ai_test_detail']:
                            test_detail = AITestDetail(
                                ai_test_case=test_case,
                                index=detail.get('index'),
                                test_step=detail.get('test_step'),
                                expected_result=detail.get('expected_result'),
                                modify_history=detail.get('modify_history'),
                                detail=detail.get('detail'),
                                remark=detail.get('remark'),
                                reserved_field0=detail.get('reserved_field0'),
                                reserved_field1=detail.get('reserved_field1'),
                                reserved_field2=detail.get('reserved_field2')
                            )
                            test_detail.save()

            return Response({'status': 'success'}, status=status.HTTP_200_OK)
        except Exception as e:
            return Response({'error': str(e)}, status=status.HTTP_400_BAD_REQUEST)